import torch
from torch import nn
from d2l import torch as d2l
from data_generate import load_data_fashion_mnist
import sys
sys.path.append('E:\AI\DL')
from Train_ import train_func as tf

alpha = 0.1
batch_size = 256
num_workers = 4
train_iter,test_iter = load_data_fashion_mnist(batch_size,num_workers)

net = nn.Sequential(nn.Flatten(),nn.Linear(784,10))
def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight,std=0.01)
net.apply(init_weights)
loss_func = nn.CrossEntropyLoss()
trainer = torch.optim.SGD(net.parameters(),alpha)
num_epochs = 10
tf.train(net,train_iter,test_iter,loss_func,num_epochs,trainer)
